"""
任务仓库

提供处理任务相关数据的专门访问方法。
"""

from typing import List, Optional, Dict, Any
from datetime import datetime, timedelta

from sqlalchemy.orm import joinedload
from sqlalchemy import and_, or_, desc, asc, func

from src.core.di import Injectable, Inject
from src.database.repository import BaseRepository
from src.models import ProcessingTask, Video
from src.models.task import TaskStatus, TaskType


@Injectable(scope="singleton")
class TaskRepository(BaseRepository[ProcessingTask]):
    """
    任务仓库
    
    提供处理任务数据的专门访问方法和业务查询。
    """
    
    def __init__(self):
        super().__init__(ProcessingTask)
    
    def get_with_video(self, task_id: int) -> Optional[ProcessingTask]:
        """获取包含视频信息的任务"""
        session = self.get_session()
        try:
            return session.query(ProcessingTask).options(
                joinedload(ProcessingTask.video)
            ).filter(ProcessingTask.id == task_id).first()
        finally:
            session.close()
    
    def find_by_status(self, status: TaskStatus) -> List[ProcessingTask]:
        """根据状态查找任务"""
        return self.get_by_field("status", status)
    
    def find_by_type(self, task_type: TaskType) -> List[ProcessingTask]:
        """根据类型查找任务"""
        return self.get_by_field("task_type", task_type)
    
    def find_pending(self) -> List[ProcessingTask]:
        """查找等待中的任务"""
        return self.find_by_status(TaskStatus.PENDING)
    
    def find_running(self) -> List[ProcessingTask]:
        """查找运行中的任务"""
        return self.find_by_status(TaskStatus.RUNNING)
    
    def find_completed(self) -> List[ProcessingTask]:
        """查找已完成的任务"""
        return self.find_by_status(TaskStatus.COMPLETED)
    
    def find_failed(self) -> List[ProcessingTask]:
        """查找失败的任务"""
        return self.find_by_status(TaskStatus.FAILED)
    
    def find_retryable(self) -> List[ProcessingTask]:
        """查找可重试的任务"""
        session = self.get_session()
        try:
            return session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.status == TaskStatus.FAILED,
                    ProcessingTask.retry_count < ProcessingTask.max_retries
                )
            ).all()
        finally:
            session.close()
    
    def find_by_worker(self, worker_id: str) -> List[ProcessingTask]:
        """根据工作节点查找任务"""
        return self.get_by_field("worker_id", worker_id)
    
    def find_by_video(self, video_id: int) -> List[ProcessingTask]:
        """根据视频查找任务"""
        return self.get_by_field("video_id", video_id)
    
    def find_recent(self, days: int = 7) -> List[ProcessingTask]:
        """查找最近的任务"""
        cutoff_date = datetime.utcnow() - timedelta(days=days)
        session = self.get_session()
        try:
            return session.query(ProcessingTask).filter(
                ProcessingTask.created_at >= cutoff_date
            ).order_by(desc(ProcessingTask.created_at)).all()
        finally:
            session.close()
    
    def find_long_running(self, hours: int = 1) -> List[ProcessingTask]:
        """查找长时间运行的任务"""
        cutoff_time = datetime.utcnow() - timedelta(hours=hours)
        session = self.get_session()
        try:
            return session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.status == TaskStatus.RUNNING,
                    ProcessingTask.started_at <= cutoff_time
                )
            ).all()
        finally:
            session.close()
    
    def find_stuck_tasks(self, hours: int = 2) -> List[ProcessingTask]:
        """查找卡住的任务"""
        cutoff_time = datetime.utcnow() - timedelta(hours=hours)
        session = self.get_session()
        try:
            return session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.status == TaskStatus.RUNNING,
                    ProcessingTask.started_at <= cutoff_time,
                    ProcessingTask.updated_at <= cutoff_time
                )
            ).all()
        finally:
            session.close()
    
    def get_next_pending_task(self, task_types: List[TaskType] = None) -> Optional[ProcessingTask]:
        """获取下一个待处理任务"""
        session = self.get_session()
        try:
            query = session.query(ProcessingTask).filter(
                ProcessingTask.status == TaskStatus.PENDING
            )
            
            if task_types:
                query = query.filter(ProcessingTask.task_type.in_(task_types))
            
            # 按优先级和创建时间排序
            return query.order_by(
                desc(ProcessingTask.priority),
                asc(ProcessingTask.created_at)
            ).first()
        finally:
            session.close()
    
    def get_queue_status(self) -> Dict[str, Any]:
        """获取任务队列状态"""
        session = self.get_session()
        try:
            # 按状态统计
            status_stats = session.query(
                ProcessingTask.status,
                func.count(ProcessingTask.id).label('count')
            ).group_by(ProcessingTask.status).all()
            
            # 按类型统计
            type_stats = session.query(
                ProcessingTask.task_type,
                func.count(ProcessingTask.id).label('count')
            ).group_by(ProcessingTask.task_type).all()
            
            # 按优先级统计
            priority_stats = session.query(
                ProcessingTask.priority,
                func.count(ProcessingTask.id).label('count')
            ).group_by(ProcessingTask.priority).all()
            
            return {
                'status_distribution': {
                    status.value: count for status, count in status_stats
                },
                'type_distribution': {
                    task_type.value: count for task_type, count in type_stats
                },
                'priority_distribution': {
                    priority: count for priority, count in priority_stats
                }
            }
        finally:
            session.close()
    
    def get_performance_statistics(self) -> Dict[str, Any]:
        """获取性能统计信息"""
        session = self.get_session()
        try:
            # 完成任务的平均处理时间
            completed_tasks = session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.status == TaskStatus.COMPLETED,
                    ProcessingTask.actual_duration.isnot(None)
                )
            ).all()
            
            if not completed_tasks:
                return {
                    'total_tasks': 0,
                    'total_completed': 0,
                    'total_failed': 0,
                    'avg_processing_time': 0,
                    'min_processing_time': 0,
                    'max_processing_time': 0,
                    'success_rate': 0
                }
            
            durations = [task.actual_duration for task in completed_tasks]
            
            # 总体统计
            total_tasks = self.count()
            completed_count = len(completed_tasks)
            failed_count = self.count(status=TaskStatus.FAILED)
            
            return {
                'total_tasks': total_tasks,
                'total_completed': completed_count,
                'total_failed': failed_count,
                'success_rate': completed_count / total_tasks if total_tasks > 0 else 0,
                'avg_processing_time': sum(durations) / len(durations),
                'min_processing_time': min(durations),
                'max_processing_time': max(durations)
            }
        finally:
            session.close()
    
    def get_worker_statistics(self) -> Dict[str, Any]:
        """获取工作节点统计信息"""
        session = self.get_session()
        try:
            worker_stats = session.query(
                ProcessingTask.worker_id,
                func.count(ProcessingTask.id).label('total_tasks'),
                func.sum(
                    func.case(
                        (ProcessingTask.status == TaskStatus.COMPLETED, 1),
                        else_=0
                    )
                ).label('completed_tasks'),
                func.avg(ProcessingTask.actual_duration).label('avg_duration')
            ).filter(
                ProcessingTask.worker_id.isnot(None)
            ).group_by(ProcessingTask.worker_id).all()
            
            stats = {}
            for worker_id, total, completed, avg_duration in worker_stats:
                stats[worker_id] = {
                    'total_tasks': total,
                    'completed_tasks': completed or 0,
                    'success_rate': (completed or 0) / total if total > 0 else 0,
                    'avg_processing_time': float(avg_duration or 0)
                }
            
            return stats
        finally:
            session.close()
    
    def cleanup_old_tasks(self, days: int = 30) -> int:
        """清理旧任务"""
        cutoff_date = datetime.utcnow() - timedelta(days=days)
        session = self.get_session()
        try:
            # 只删除已完成或失败的旧任务
            old_tasks_count = session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.completed_at <= cutoff_date,
                    or_(
                        ProcessingTask.status == TaskStatus.COMPLETED,
                        ProcessingTask.status == TaskStatus.FAILED,
                        ProcessingTask.status == TaskStatus.CANCELLED
                    )
                )
            ).count()
            
            session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.completed_at <= cutoff_date,
                    or_(
                        ProcessingTask.status == TaskStatus.COMPLETED,
                        ProcessingTask.status == TaskStatus.FAILED,
                        ProcessingTask.status == TaskStatus.CANCELLED
                    )
                )
            ).delete(synchronize_session=False)
            
            session.commit()
            return old_tasks_count
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
    
    def cancel_stuck_tasks(self, hours: int = 2) -> int:
        """取消卡住的任务"""
        stuck_tasks = self.find_stuck_tasks(hours)
        
        cancelled_count = 0
        for task in stuck_tasks:
            task.cancel()
            self.update_from_dict(task.id, task.to_dict())
            cancelled_count += 1
        
        return cancelled_count
    
    def retry_failed_tasks(self, max_retries: int = None) -> int:
        """重试失败的任务"""
        retryable_tasks = self.find_retryable()
        
        if max_retries:
            retryable_tasks = [
                task for task in retryable_tasks 
                if task.retry_count < max_retries
            ]
        
        retried_count = 0
        for task in retryable_tasks:
            task.retry()
            self.update_from_dict(task.id, task.to_dict())
            retried_count += 1
        
        return retried_count
    
    def get_task_history(self, video_id: int) -> List[ProcessingTask]:
        """获取视频的任务历史"""
        return self.find(
            filters={"video_id": video_id},
            order_by="created_at",
            order_desc=True
        )
    
    def get_latest_task_by_type(self, video_id: int, task_type: TaskType) -> Optional[ProcessingTask]:
        """获取视频的最新指定类型任务"""
        session = self.get_session()
        try:
            return session.query(ProcessingTask).filter(
                and_(
                    ProcessingTask.video_id == video_id,
                    ProcessingTask.task_type == task_type
                )
            ).order_by(desc(ProcessingTask.created_at)).first()
        finally:
            session.close()
